#!/usr/bin/env python
from glob import glob
from astropy.io import fits
from astropy.table import Table
import sys
import os
import numpy as np

def crop_overscan(image, imgconf):
    output = np.zeros(imgconf.mos_shape, dtype=image.dtype)
    for chan in range(imgconf.nchan):
        x1 = imgconf.raw_x0[chan]
        y1 = imgconf.raw_y0[chan]
        cimage = image[y1:y1 + imgconf.raw_ny[chan], x1:x1 + imgconf.raw_nx[chan]]

        if imgconf.raw_xflip[chan]:
            x1 = imgconf.raw_nx[chan] - imgconf.x1[chan] - imgconf.nx[chan]
        else:
            x1 = imgconf.x1[chan]
        if imgconf.raw_yflip[chan]:
            y1 = imgconf.raw_ny[chan] - imgconf.y1[chan] - imgconf.ny[chan]
        else:
            y1 = imgconf.y1[chan]
        cimage = cimage[y1:y1 + imgconf.ny[chan], x1:x1 + imgconf.nx[chan]]

        if imgconf.mos_xflip[chan]:
            cimage = cimage[:, ::-1]
        if imgconf.mos_yflip[chan]:
            cimage = cimage[::-1, :]
        x1 = imgconf.mos_x0[chan]
        y1 = imgconf.mos_y0[chan]
        output[y1:y1 + imgconf.ny[chan], x1:x1 + imgconf.nx[chan]] = cimage
    return output

class ConfHR9090BSI:
    def __init__(self):
        self.name = 'hr9090bsi'
        self.type = 'cmos'
        self.tsys = 188
        self.glass = 'none'
        self.wave = -99.0

        # channel layout
        self.nchan = 2
        self.col = [0, 1]
        self.row = [0, 0]
        self.overscan = 1

        # raw image
        self.raw_shape = (9222, 8976)
        self.raw_x0 = [0, 4352]
        self.raw_y0 = [0, 0]
        self.raw_nx = [4352, 4624]
        self.raw_ny = [9222, 9222]
        self.raw_xflip = [0, 1]  # flip for overscan region alignment
        self.raw_yflip = [0, 0]

        # exposure region, prescan and overscan (in raw-flipped channel cutout)
        self.x1 = [38, 38]  # EB 34 + dummy 4
        self.y1 = [92, 92]  # dummy 92
        self.nx = [4314, 4586]
        self.ny = [9120, 9120]
        self.overscan_x1 = [0, 0]
        self.overscan_x2 = [34, 34]  # only use EB
        self.overscan_y1 = [0, 0]
        self.overscan_y2 = [92, 92]

        # mosaic image (without prescan and overscan)
        self.mos_shape = (9120, 8900)
        self.mos_x0 = [0, 4314]
        self.mos_y0 = [0, 0]
        self.mos_xflip = [0, 0]  # flip the exposure region for mosaic
        self.mos_yflip = [0, 0]

gain_map = fits.getdata(sys.argv[1])
if not os.path.exists("emap.fits"):
    for i in range(1000):
        print(i+1)
        if i==0:
            data0 = fits.getdata("flat-0.0-"+str(i+1)+".fits")/1000
            data = crop_overscan(data0, ConfHR9090BSI())*gain_map
        else:
            data0 = fits.getdata("flat-0.0-"+str(i+1)+".fits")/1000
            data += crop_overscan(data0, ConfHR9090BSI())*gain_map
    e_map = data.copy()
    fits.writeto("emap.fits",data=e_map,overwrite=True)
else:
    e_map = fits.getdata("emap.fits")
tab = Table()
flux = []
std = []
trap_map = np.zeros(e_map.shape)
trap_times = []
for i in range(1000):
    print(i)
    data = fits.getdata("flat-0.0-"+str(i+1)+".fits")
    data = crop_overscan(data, ConfHR9090BSI())
    print(data.shape)
    e_map1 = data*gain_map
    flux.append(np.mean(e_map1))
    std.append(np.std(e_map1))
    trap_map += (abs(e_map-e_map1)>200).astype('int')
    trap_times.append(np.sum((abs(e_map-e_map1)>200).astype('int')))
np.savetxt("trap_times.txt",np.array(trap_times))
fits.writeto("trap.fits",data=trap_map,overwrite=True)
print("trap:",len(np.where(trap_map>1)[0])/81095936)
tab['flux'] = np.array(flux)
tab['std'] = np.array(std)
tab.write("trap.tab",format='ipac',overwrite=True)
